# -*- coding: utf-8 -*-
"""
@Auther: nieyuwei
@Create Time: 2024/01/25
"""

from PIL import Image
from skimage import io as skio  # 用于图像的IO和变换

from scripts.kp.data_process import *

transform_img = transforms.Compose([
    transforms.ToTensor(),  # 将图像(Image)转成Tensor
])

# model = torch.load('/alidata/work/models/KPDM/train_top_1706088268/epoch_16.pth')
model_path = os.getenv('MODELSCOPE_CACHE_KP', "d:/tmp/work/pth/epoch_16.pth")
model = torch.load(model_path)
# model = torch.load('/root/.cache/kf/epoch_16.pth', map_location=torch.device('cpu'))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def distance(L, R):
    return ((L[0] - R[0]) ** 2 + (L[1] - R[1]) ** 2) ** 0.5


def cal_pos(L, R):
    a = -0.64
    b = -0.87
    d = distance(L, R)
    cy = (L[1] + R[1]) / 2 if abs(L[1] - R[1]) < d * 0.3 else max(L[1], R[1])
    x = int(round((L[0] + R[0]) / 2 + a * d, 0))
    y = int(round(cy + b * d, 0))
    s = int(round(1.32 * d, 0))
    return x, y, s


def merge(background, foreground, x, y, s):
    if background.mode != 'RGBA':
        background = background.convert('RGBA')
    # 创建一个与背景图像相同大小的透明图层
    combined_image = Image.new('RGBA', background.size, (0, 0, 0, 0))
    foreground = foreground.resize((s, s))
    # 将前景图叠加在透明图层上
    combined_image.paste(foreground, (x, y), mask=foreground)

    # 将透明图层叠加在背景图上
    final_image = Image.alpha_composite(background, combined_image)
    return final_image


def kp_infer(image_file):
    image = skio.imread(image_file)
    h, w, c = image.shape
    image = image[:, :, :3]
    img = transform.resize(image, (64, 64))
    img = transform_img(img) / 255.0
    img = img.to(device).float()
    coords, _ = model(img.unsqueeze(0))
    coords = coords.squeeze().cpu().detach().numpy()
    coords = coords * 32.0 + 31.5
    coords[:, 0] = coords[:, 0] * w / 64
    coords[:, 1] = coords[:, 1] * h / 64
    xl = min(coords[2, 0], coords[4, 0])
    yl = min(max(coords[2, 1], coords[0, 1]), coords[4, 1])
    xr = max(coords[3, 0], coords[5, 0])
    yr = min(max(coords[3, 1], coords[1, 1]), coords[5, 1])

    return [xl, yl], [xr, yr]
